import pandas as pd
import matplotlib.pyplot as plt
from utils import Config
from data_analysis.correlation_analysis import calculate_column_correlation

import matplotlib.pyplot as plt
import seaborn as sns

plt.rcParams['font.sans-serif'] = ['SimHei']  # 解决中文乱码问题
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题


def read_excel():
    if Config.sheet_name == "不同年月气象信息":
        df = pd.read_excel(Config.excel_data_path, sheet_name=Config.sheet_name, header=[0, 1])
        return df
    elif Config.sheet_name == "不同年份人口经济":
        df = pd.read_excel(Config.excel_data_path, sheet_name=Config.sheet_name, header=[0, 1])
        return df
    elif Config.sheet_name == "不同年份土地利用类型":
        df = pd.read_excel(Config.excel_data_path, sheet_name=Config.sheet_name, header=[0, 1])
        return df
    elif Config.sheet_name == "地理基础信息":
        df = pd.read_excel(Config.excel_data_path, sheet_name=Config.sheet_name, header=[0, 1])
        return df
    else:
        return None


def read_excel_sheet(sheet_name):
    df = pd.read_excel(Config.excel_data_path, sheet_name=sheet_name, header=[0, 1])
    return df


def select_data_by_year(df):
    # 按照年的范围选择数据
    # 删除第一层的列名，使其变为单层索引
    if Config.sheet_name == "不同年月气象信息":
        df.columns = ['流域网格编码', '年份', '月份', '降水量（mm）']
        filtered_df = df[df['年份'].isin(Config.year_range)]
        return filtered_df
    elif Config.sheet_name == "不同年份人口经济":
        df.columns = ['流域网格编码', '年份', 'GDP（万元）', '人口数量（人）']
        filtered_df = df[df['年份'].isin(Config.year_range)]
        return filtered_df
    elif Config.sheet_name == "不同年份土地利用类型":
        df.columns = ['流域网格编码', '年份', '农田（km2）', '森林（km2）', '灌木（km2）', '草原（km2）',
                      '水域（km2）', '冰雪（km2）', '裸地（km2）', '不透水面（km2）', '湿地（km2）']
        filtered_df = df[df['年份'].isin(Config.year_range)]
        return filtered_df
    elif Config.sheet_name == "地理基础信息":
        df.columns = ['流域网格编码', '单个网格流域面积（km2) ', '经度（°）', '纬度（°）', '坡度（均值/°）', '坡向（均值/°）']
        filtered_df = df[df['年份'].isin(Config.year_range)]
        return filtered_df
    else:
        return None


def select_data_by_month(df):
    # 按照月的范围选择数据
    start_month, end_month = Config.month_range.split("-")
    month_range = range(int(start_month), int(end_month) + 1)
    if Config.sheet_name == "不同年月气象信息":
        df.columns = ['流域网格编码', '年份', '月份', '降水量（mm）']
        filtered_df = df[df['月份'].isin(month_range)]
        return filtered_df
    elif Config.sheet_name == "不同年份人口经济":
        df.columns = ['流域网格编码', '年份', 'GDP（万元）', '人口数量（人）']
        filtered_df = df[df['月份'].isin(month_range)]
        return filtered_df
    else:
        return None


def select_data_by_gridcell(df):
    start_grid, end_grid = Config.gridcell_range.split("-")
    print(start_grid, end_grid)
    grid_range = range(int(start_grid), int(end_grid) + 1)
    if Config.sheet_name == "不同年月气象信息":
        df.columns = ['流域网格编码', '年份', '月份', '降水量（mm）']
        filtered_df = df[df['流域网格编码'].isin(grid_range)]
        return filtered_df
    elif Config.sheet_name == "不同年份人口经济":
        df.columns = ['流域网格编码', '年份', 'GDP（万元）', '人口数量（人）']
        filtered_df = df[df['流域网格编码'].isin(grid_range)]
        return filtered_df
    elif Config.sheet_name == "不同年份土地利用类型":
        df.columns = ['流域网格编码', '年份', '农田（km2）', '森林（km2）', '灌木（km2）', '草原（km2）',
                      '水域（km2）', '冰雪（km2）', '裸地（km2）', '不透水面（km2）', '湿地（km2）']
        filtered_df = df[df['流域网格编码'].isin(grid_range)]
        return filtered_df
    elif Config.sheet_name == "地理基础信息":
        df.columns = ['流域网格编码', '单个网格流域面积（km2) ', '经度（°）', '纬度（°）', '坡度（均值/°）', '坡向（均值/°）']
        filtered_df = df[df['流域网格编码'].isin(grid_range)]
        return filtered_df


def calculate_monthly_average_by_year(df):
    """
    根据年份和月份对降水量进行分组并计算平均值。

    参数:
    df - 包含年份、月份、降水量的 DataFrame

    返回:
    按年和月份分组后的降水量平均值 DataFrame
    """
    # 按照年份和月份进行分组，然后计算降水量的均值
    if Config.sheet_name == "不同年月气象信息":
        monthly_avg_by_year = df.groupby(['年份', '月份', '流域网格编码'])['降水量（mm）'].mean().reset_index()
        return monthly_avg_by_year
    else:
        return None


def calculate_yearly_sum_by_year(df):
    # 按照年份和月份进行分组，然后计算降水量的均值
    if Config.sheet_name == "不同年月气象信息":
        monthly_avg_by_year = df.groupby(['年份', '流域网格编码'])['降水量（mm）'].sum().reset_index()
        return monthly_avg_by_year
    else:
        return None


def calculate_total_persons_average_by_year(df):
    """
    根据年份和月份对降水量进行分组并计算平均值。

    参数:
    df - 包含年份、月份、降水量的 DataFrame

    返回:
    按年和月份分组后的降水量平均值 DataFrame
    """
    # 按照年份和月份进行分组，然后计算降水量的均值
    if Config.sheet_name == "不同年份人口经济":
        persons_avg_by_year = df.groupby('年份')['人口数量（人）'].sum().reset_index()
        return persons_avg_by_year
    else:
        return None


def calculate_total_gdp_average_by_year(df):
    # 按照年份和月份进行分组，然后计算降水量的均值
    if Config.sheet_name == "不同年份人口经济":
        gdp_avg_by_year = df.groupby('年份')['GDP（万元）'].sum().reset_index()
        return gdp_avg_by_year
    else:
        return None


def calculate_total_landuse_average_by_year(df):
    if Config.sheet_name == "不同年份土地利用类型":
        landuse_sum_by_year = df.groupby('年份')[['农田（km2）', '森林（km2）', '灌木（km2）', '草原（km2）',
                                                '水域（km2）', '冰雪（km2）', '裸地（km2）', '不透水面（km2）',
                                                '湿地（km2）']].sum().reset_index()
        return landuse_sum_by_year
    else:
        return None


def plot_over_time(df):
    """
    绘制基于时间（年份+月份）的降水量趋势图。

    参数:
    df - 包含 '年份', '月份' 和 '降水量（mm）' 列的 DataFrame
    """
    # 创建新的时间列，格式为 YYYY-MM
    if Config.sheet_name == "不同年月气象信息":
        df['时间'] = df['年份'].astype(str) + '-' + df['月份'].astype(str).str.zfill(2) + '-' + df['流域网格编码'].astype(
            str).str.zfill(3)

        # 绘制降水量趋势图
        plt.figure(figsize=(48, 12))
        plt.plot(df['时间'], df['降水量（mm）'], marker='o', linestyle='-', color='b')
        plt.title('Precipitation Over Time', fontsize=16)
        plt.xlabel('Time (Year-Month)', fontsize=12)
        plt.ylabel('Average Precipitation (mm)', fontsize=12)
        plt.xticks(rotation=45, ha='right')  # 旋转x轴刻度标签，使其可读
        plt.grid(True)
        plt.tight_layout()  # 自动调整图像，避免标签重叠
        plt.show()
    elif Config.sheet_name == "不同年份人口经济":
        # 按年绘制人口变化
        df['时间'] = df['年份'].astype(str) + '-' + df['流域网格编码'].astype(
            str).str.zfill(3)
        if Config.plot_gdp:
            plt.figure(figsize=(12, 6))
            plt.plot(df['时间'], df['GDP（万元）'], marker='o', linestyle='-', color='b')
            plt.title('GDP changes Over Time', fontsize=16)
            plt.xlabel('Time (Year)', fontsize=12)
            plt.ylabel('Average Precipitation (mm)', fontsize=12)
            plt.xticks(rotation=45, ha='right')  # 旋转x轴刻度标签，使其可读
            plt.grid(True)
            plt.tight_layout()  # 自动调整图像，避免标签重叠
            plt.show()
        else:
            plt.figure(figsize=(12, 6))
            plt.plot(df['时间'], df['人口数量（人）'], marker='o', linestyle='-', color='b')
            plt.title('Demographic changes Over Time', fontsize=16)
            plt.xlabel('Time (Year)', fontsize=12)
            plt.ylabel('Average Precipitation (mm)', fontsize=12)
            plt.xticks(rotation=45, ha='right')  # 旋转x轴刻度标签，使其可读
            plt.grid(True)
            plt.tight_layout()  # 自动调整图像，避免标签重叠
            plt.show()
    elif Config.sheet_name == "不同年份土地利用类型":
        df['时间'] = df['年份'].astype(str)
        plt.figure(figsize=(12, 6))
        # 绘制每种土地利用类型的变化曲线
        plt.plot(df['时间'], df['农田（km2）'], marker='o', linestyle='-', label='Farmland (km2)')
        plt.plot(df['时间'], df['森林（km2）'], marker='o', linestyle='-', label='Forest (km2)')
        plt.plot(df['时间'], df['灌木（km2）'], marker='o', linestyle='-', label='Shrubland (km2)')
        plt.plot(df['时间'], df['草原（km2）'], marker='o', linestyle='-', label='Grassland (km2)')
        plt.plot(df['时间'], df['水域（km2）'], marker='o', linestyle='-', label='Water (km2)')
        plt.plot(df['时间'], df['冰雪（km2）'], marker='o', linestyle='-', label='Snow/Ice (km2)')
        plt.plot(df['时间'], df['裸地（km2）'], marker='o', linestyle='-', label='Bare land (km2)')
        plt.plot(df['时间'], df['不透水面（km2）'], marker='o', linestyle='-', label='Impervious surface (km2)')
        plt.plot(df['时间'], df['湿地（km2）'], marker='o', linestyle='-', label='Wetland (km2)')

        # 设置图表标题和标签
        plt.title('Land Use Changes Over Time', fontsize=16)
        plt.xlabel('Time (Year)', fontsize=12)
        plt.ylabel('Area (km2)', fontsize=12)
        plt.xticks(rotation=45, ha='right')  # 旋转x轴刻度标签，使其可读
        plt.grid(True)
        plt.legend()  # 显示图例
        plt.tight_layout()  # 自动调整图像，避免标签重叠
        plt.show()

    elif Config.sheet_name == "地理基础信息":
        plt.figure(figsize=(12, 6))
        plt.plot(df['流域网格编码'], df['单个网格流域面积（km2) '], marker='o', linestyle='-', label='单个网格流域面积（km2) ')
        plt.plot(df['流域网格编码'], df['经度（°）'], marker='o', linestyle='-', label='经度（°）')
        plt.plot(df['流域网格编码'], df['纬度（°）'], marker='o', linestyle='-', label='纬度（°）')
        plt.plot(df['流域网格编码'], df['坡度（均值/°）'], marker='o', linestyle='-', label='坡度（均值/°）')
        plt.plot(df['流域网格编码'], df['坡向（均值/°）'], marker='o', linestyle='-', label='坡向（均值/°）')
        plt.title('地理基础信息', fontsize=16)
        plt.xlabel('网格编码', fontsize=12)
        plt.ylabel('基础信息', fontsize=12)
        plt.xticks(rotation=45, ha='right')  # 旋转x轴刻度标签，使其可读
        plt.grid(True)
        plt.legend()  # 显示图例
        plt.tight_layout()  # 自动调整图像，避免标签重叠
        plt.show()


def filter_invalid_gdp(df):
    """
    过滤掉 GDP（万元） 列中包含 '——' 的数据。

    参数:
    df - 包含 GDP 列的 DataFrame

    返回:
    过滤后的 DataFrame
    """
    # 使用条件过滤掉 'GDP（万元）' 列中包含 '——' 的行
    filtered_df = df[df['GDP（万元）'] != '——']

    return filtered_df


def calculate_gdpandpopulation_change_rate_with_grid(df):
    # 按年份排序并提取流域网格数据
    df = df.sort_values(by='年份')

    # 提取土地利用的列名（即从第3列开始的列是土地利用类型）
    land_use_columns = df.columns[2:]

    # 获取网格编号（假设每个网格都有唯一编号，且编号保持一致）
    grid_id = df['流域网格编码'].unique()[0]  # 获取唯一的网格编号

    # 创建一个空的列表用于存储变化速率
    results = []

    # 获取年份列表
    years = df['年份'].unique()

    # 遍历每种土地利用类型，计算每个时间段的变化速率
    for land_use_type in land_use_columns:
        change_rate = pd.DataFrame()  # 临时DataFrame存储当前土地利用类型的变化速率
        for i in range(1, len(years)):
            year1 = years[i - 1]
            year2 = years[i]
            # 获取指定年份的数据
            value_year1 = df[df['年份'] == year1][land_use_type].values[0]
            value_year2 = df[df['年份'] == year2][land_use_type].values[0]

            # 计算变化速率
            rate = (value_year2 - value_year1) / (year2 - year1)
            change_rate[f'Change_rate_{year1}_{year2}'] = [rate]

        # 找到最大变化速率和对应的时间段
        max_change = change_rate.max(axis=1).values[0]
        max_change_time = change_rate.idxmax(axis=1).values[0]

        # 创建结果DataFrame，添加网格编号和土地利用类型
        results.append(pd.DataFrame({
            'Max_change_rate': [max_change],
            'Max_change_time': [max_change_time],
            'GDP和人口': [land_use_type],
            '网格编码': [grid_id]
        }))

    # 返回所有土地利用类型的拼接结果
    return pd.concat(results, ignore_index=True)


def calculate_land_use_change_rate_with_grid(df):
    # 按年份排序并提取流域网格数据
    df = df.sort_values(by='年份')

    # 提取土地利用的列名（即从第3列开始的列是土地利用类型）
    land_use_columns = df.columns[2:]

    # 获取网格编号（假设每个网格都有唯一编号，且编号保持一致）
    grid_id = df['流域网格编码'].unique()[0]  # 获取唯一的网格编号

    # 创建一个空的列表用于存储变化速率
    results = []

    # 获取年份列表
    years = df['年份'].unique()

    # 遍历每种土地利用类型，计算每个时间段的变化速率
    for land_use_type in land_use_columns:
        change_rate = pd.DataFrame()  # 临时DataFrame存储当前土地利用类型的变化速率
        for i in range(1, len(years)):
            year1 = years[i - 1]
            year2 = years[i]
            # 获取指定年份的数据
            value_year1 = df[df['年份'] == year1][land_use_type].values[0]
            value_year2 = df[df['年份'] == year2][land_use_type].values[0]

            # 计算变化速率
            rate = (value_year2 - value_year1) / (year2 - year1)
            change_rate[f'Change_rate_{year1}_{year2}'] = [rate]

        # 找到最大变化速率和对应的时间段
        max_change = change_rate.max(axis=1).values[0]
        max_change_time = change_rate.idxmax(axis=1).values[0]

        # 创建结果DataFrame，添加网格编号和土地利用类型
        results.append(pd.DataFrame({
            'Max_change_rate': [max_change],
            'Max_change_time': [max_change_time],
            '土地利用类型': [land_use_type],
            '网格编码': [grid_id]
        }))

    # 返回所有土地利用类型的拼接结果
    return pd.concat(results, ignore_index=True)


def calculate_rain_change_rate_with_grid(df):
    # 按年份排序并提取流域网格数据
    df = df.sort_values(by='年份')

    rain_columns = df.columns[2:]

    # 获取网格编号（假设每个网格都有唯一编号，且编号保持一致）
    grid_id = df['流域网格编码'].unique()[0]  # 获取唯一的网格编号

    # 创建一个空的列表用于存储变化速率
    results = []

    # 获取年份列表
    years = df['年份'].unique()

    change_rate = pd.DataFrame()
    for i in range(1, len(years)):
        year1 = years[i - 1]
        year2 = years[i]
        value_year1 = df[df['年份'] == year1][rain_columns].values[0]
        value_year2 = df[df['年份'] == year2][rain_columns].values[0]

        # 计算变化速率
        rate = (value_year2 - value_year1) / (year2 - year1)
        change_rate[f'Change_rate_{year1}_{year2}'] = [rate]

    max_change = change_rate.max(axis=1).values[0]
    max_change_time = change_rate.idxmax(axis=1).values[0]

    results.append(pd.DataFrame({
        'Max_change_rate': max_change,
        'Max_change_time': [max_change_time],
        '土地利用类型': rain_columns,
        '网格编码': [grid_id]
    }))

    # 返回所有土地利用类型的拼接结果
    return pd.concat(results, ignore_index=True)


def main():
    if Config.self_analysis:
        if Config.sheet_name == "不同年月气象信息":
            df = read_excel()
            df = select_data_by_year(df)
            df = select_data_by_month(df)
            df = select_data_by_gridcell(df)
            # monthly_avg_by_year = calculate_monthly_average_by_year(df)
            plot_over_time(df)
        elif Config.sheet_name == "不同年份人口经济":
            df = read_excel()
            df = select_data_by_year(df)
            df = select_data_by_gridcell(df)
            if Config.plot_gdp == True:
                filtered_gdp_df = filter_invalid_gdp(df)
                print(filtered_gdp_df.shape[0])
                # total_gdp_avg_by_year = calculate_total_gdp_average_by_year(filtered_gdp_df)
                plot_over_time(filtered_gdp_df)
            else:
                # total_persons_avg_by_year = calculate_total_persons_average_by_year(df)
                plot_over_time(df)
        elif Config.sheet_name == "不同年份土地利用类型":
            df = read_excel()
            df = select_data_by_year(df)
            df = select_data_by_gridcell(df)
            total_landuse_average_by_year = calculate_total_landuse_average_by_year(df)
            plot_over_time(total_landuse_average_by_year)
        elif Config.sheet_name == "地理基础信息":
            df = read_excel()
            df = select_data_by_gridcell(df)
            plot_over_time(df)

    if Config.Change_analysis:
        if Config.sheet_name == "不同年份土地利用类型":
            all_results = []  # 用于存储每个网格的结果
            df = read_excel()
            df = select_data_by_year(df)
            old_cell_range = Config.gridcell_range
            for cell_number in range(1, 324):
                Config.gridcell_range = str(cell_number) + "-" + str(cell_number)
                df_cell = select_data_by_gridcell(df)
                change_rate_df = calculate_land_use_change_rate_with_grid(df_cell)
                all_results.append(change_rate_df)
            Config.gridcell_range = old_cell_range
            # 拼接所有结果
            final_df = pd.concat(all_results, axis=0)
            # 保存结果到CSV文件
            final_df.to_csv(Config.output_land_use_change_rate, index=False, encoding="utf-8")
        elif Config.sheet_name == "不同年月气象信息":
            all_results = []  # 用于存储每个网格的结果
            df = read_excel()
            df = select_data_by_year(df)
            df = select_data_by_month(df)
            old_cell_range = Config.gridcell_range
            for cell_number in range(1, 324):
                Config.gridcell_range = str(cell_number) + "-" + str(cell_number)
                df_cell = select_data_by_gridcell(df)
                yearly_sum_by_year = calculate_yearly_sum_by_year(df_cell)
                change_rate_df = calculate_rain_change_rate_with_grid(yearly_sum_by_year)
                all_results.append(change_rate_df)
            Config.gridcell_range = old_cell_range
            # 拼接所有结果
            final_df = pd.concat(all_results, axis=0)
            # 保存结果到CSV文件
            final_df.to_csv(Config.rain_use_change_rate, index=False, encoding="utf-8")
        elif Config.sheet_name == "不同年份人口经济":
            all_results = []
            df = read_excel()
            df = select_data_by_year(df)
            old_cell_range = Config.gridcell_range
            for cell_number in range(1, 324):
                Config.gridcell_range = str(cell_number) + "-" + str(cell_number)
                df_cell = select_data_by_gridcell(df)
                change_rate_df = calculate_gdpandpopulation_change_rate_with_grid(df_cell)
                all_results.append(change_rate_df)
            Config.gridcell_range = old_cell_range
            # 拼接所有结果
            final_df = pd.concat(all_results, axis=0)
            # 保存结果到CSV文件
            final_df.to_csv(Config.gdp_population_change_rate, index=False, encoding="utf-8")
    # 相关性分析
    if Config.corr_analysis:
        df1 = read_excel_sheet("不同年月气象信息")
        old_sheet_name = Config.sheet_name
        Config.sheet_name = "不同年月气象信息"
        df1 = select_data_by_year(df1)
        df1 = select_data_by_gridcell(df1)
        df1 = calculate_yearly_sum_by_year(df1)  # 由于其余表没有月份 所以需要把按年累计降雨量

        df2 = read_excel_sheet("不同年份人口经济")
        Config.sheet_name = "不同年份人口经济"
        df2 = select_data_by_year(df2)
        df2 = select_data_by_gridcell(df2)

        df3 = read_excel_sheet("不同年份土地利用类型")
        Config.sheet_name = "不同年份土地利用类型"
        df3 = select_data_by_year(df3)
        df3 = select_data_by_gridcell(df3)

        df4 = read_excel_sheet("地理基础信息")
        Config.sheet_name = "地理基础信息"
        df4 = select_data_by_gridcell(df4)
        Config.sheet_name = old_sheet_name

        # 合并数据，使用“流域网格编码”和“年份”作为键
        merged_df = pd.merge(df1, df2, on=["流域网格编码", "年份"], how="inner")
        merged_df = pd.merge(merged_df, df3, on=["流域网格编码", "年份"], how="inner")
        merged_df = pd.merge(merged_df, df4, on=["流域网格编码"], how="inner")
        merged_df.fillna(0, inplace=True)  # 可以将缺失值填充为0，或选择其他合理的填充值
        # 选择需要进行相关性分析的列，例如气温、GDP、人口数量、土地利用类型等
        columns_of_interest = Config.columns_of_interest
        correlation_df = merged_df[columns_of_interest]

        # 计算相关性
        correlation_matrix = correlation_df.corr(method='pearson')

        # 输出相关性矩阵
        print("相关性矩阵：")
        print(correlation_matrix)

        plt.figure(figsize=(10, 8))
        sns.heatmap(correlation_matrix, annot=True, cmap="coolwarm", linewidths=0.5, fmt=".2f", annot_kws={"size": 12})
        plt.title("不同变量之间的相关性矩阵")
        plt.show()


if __name__ == "__main__":
    main()
